package dycaps

import (
	"crypto/rand"
	"fmt"
	"math/big"

	"github.com/consensys/gnark-crypto/ecc/bn254"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr"
)

// GnarkKZG 基于gnark-crypto的KZG实现（使用BN254曲线）
type GnarkKZG struct {
	// 可信设置参数 g^{tau^i} for i in 0..maxDegree
	g1Powers   []bn254.G1Affine  // Powers of tau in G1
	g2Powers   []bn254.G2Affine  // Powers of tau in G2 (only need first 2)
	maxDegree  uint64
	tau        *fr.Element       // Secret for testing (production should use MPC)
}

// NewGnarkKZG 创建新的GnarkKZG实例
func NewGnarkKZG(maxDegree uint64) (*GnarkKZG, error) {
	// 生成可信设置 (在生产环境中应该使用MPC生成的可信设置)
	// 这里使用固定的tau用于测试
	var tau fr.Element
	tau.SetRandom() // 在生产环境中，tau应该被销毁

	kzg := &GnarkKZG{
		g1Powers:  make([]bn254.G1Affine, maxDegree+1),
		g2Powers:  make([]bn254.G2Affine, 2),
		maxDegree: maxDegree,
		tau:       &tau,
	}

	// 生成G1和G2的powers of tau
	if err := kzg.setupTrustedSetup(); err != nil {
		return nil, fmt.Errorf("failed to setup trusted setup: %w", err)
	}

	return kzg, nil
}

// setupTrustedSetup 生成可信设置
func (k *GnarkKZG) setupTrustedSetup() error {
	// G1生成元
	var g1Gen bn254.G1Affine
	g1Gen.X.SetString("1")
	g1Gen.Y.SetString("2")

	// G2生成元
	var g2Gen bn254.G2Affine
	// BN254 G2生成元的坐标
	g2Gen.X.SetString("10857046999023057135944570762232829481370756359578518086990519993285655852781",
		"11559732032986387107991004021392285783925812861821192530917403151452391805634")
	g2Gen.Y.SetString("8495653923123431417604973247489272438418190587263600148770280649306958101930",
		"4082367875863433681332203403145435568316851327593401208105741076214120093531")

	// 计算 g1^{tau^i}
	var tauPower fr.Element
	tauPower.SetOne()

	for i := uint64(0); i <= k.maxDegree; i++ {
		var scalar big.Int
		tauPower.BigInt(&scalar)
		k.g1Powers[i].ScalarMultiplication(&g1Gen, &scalar)
		
		if i < k.maxDegree {
			tauPower.Mul(&tauPower, k.tau)
		}
	}

	// 计算 g2^1 和 g2^tau
	k.g2Powers[0] = g2Gen
	var tauBig big.Int
	k.tau.BigInt(&tauBig)
	k.g2Powers[1].ScalarMultiplication(&g2Gen, &tauBig)

	return nil
}

// Commit 提交多项式并返回承诺
func (k *GnarkKZG) Commit(coefficients []byte) ([]byte, error) {
	// 将字节转换为fr.Element切片
	numCoeffs := len(coefficients) / 32 // 假设每个系数32字节
	if numCoeffs == 0 {
		numCoeffs = 1
	}
	
	polynomial := make([]fr.Element, numCoeffs)
	for i := 0; i < numCoeffs && i*32 < len(coefficients); i++ {
		end := (i + 1) * 32
		if end > len(coefficients) {
			end = len(coefficients)
		}
		polynomial[i].SetBytes(coefficients[i*32 : end])
	}

	// 计算承诺 C = sum(coeff_i * g1^{tau^i})
	var commitment bn254.G1Affine
	commitment.SetInfinity()

	for i := 0; i < numCoeffs && i <= int(k.maxDegree); i++ {
		var scalar big.Int
		polynomial[i].BigInt(&scalar)
		
		var term bn254.G1Affine
		term.ScalarMultiplication(&k.g1Powers[i], &scalar)
		commitment.Add(&commitment, &term)
	}

	// 序列化承诺点
	commitmentBytes := commitment.Bytes()
	return commitmentBytes[:], nil
}

// CreateWitness 为指定位置创建KZG证明
func (k *GnarkKZG) CreateWitness(coefficients []byte, index int, value []byte) ([]byte, error) {
	// 将字节转换为fr.Element切片
	numCoeffs := len(coefficients) / 32
	if numCoeffs == 0 {
		numCoeffs = 1
	}
	
	polynomial := make([]fr.Element, numCoeffs)
	for i := 0; i < numCoeffs && i*32 < len(coefficients); i++ {
		end := (i + 1) * 32
		if end > len(coefficients) {
			end = len(coefficients)
		}
		polynomial[i].SetBytes(coefficients[i*32 : end])
	}

	// 评估点
	var z fr.Element
	z.SetInt64(int64(index))

	// 计算 f(z)
	var y fr.Element
	evaluatePolynomialFr(polynomial, &z, &y)

	// 计算商多项式 q(x) = (f(x) - f(z)) / (x - z)
	quotient := make([]fr.Element, len(polynomial))
	for i := range quotient {
		quotient[i].Set(&polynomial[i])
	}
	
	// 减去 f(z)
	quotient[0].Sub(&quotient[0], &y)
	
	// 除以 (x - z) - 使用综合除法
	for i := len(quotient) - 1; i > 0; i-- {
		quotient[i-1].Add(&quotient[i-1], &fr.Element{}).Mul(&quotient[i], &z)
	}

	// 计算证明 pi = sum(q_i * g1^{tau^i})
	var proof bn254.G1Affine
	proof.SetInfinity()

	for i := 0; i < len(quotient)-1 && i <= int(k.maxDegree); i++ {
		var scalar big.Int
		quotient[i].BigInt(&scalar)
		
		var term bn254.G1Affine
		term.ScalarMultiplication(&k.g1Powers[i], &scalar)
		proof.Add(&proof, &term)
	}

	// 序列化证明
	proofBytes := proof.Bytes()
	return proofBytes[:], nil
}

// VerifyEval 验证KZG评估证明  
func (k *GnarkKZG) VerifyEval(commitmentBytes []byte, index int, valueBytes []byte, witnessBytes []byte) bool {
	// 反序列化承诺
	var commitment bn254.G1Affine
	if _, err := commitment.SetBytes(commitmentBytes); err != nil {
		return false
	}

	// 反序列化证明
	var proof bn254.G1Affine
	if _, err := proof.SetBytes(witnessBytes); err != nil {
		return false
	}

	// 评估点
	var z fr.Element
	z.SetInt64(int64(index))

	// 评估值
	var y fr.Element
	if len(valueBytes) > 32 {
		valueBytes = valueBytes[:32]
	}
	y.SetBytes(valueBytes)

	// 验证配对等式: e(C - y*G1, G2) = e(pi, tau*G2 - z*G2)
	// 等价于: e(C - y*G1, G2) * e(-pi, tau*G2 - z*G2) = 1
	
	// 计算 C - y*G1
	var yG1 bn254.G1Affine
	var yBig big.Int
	y.BigInt(&yBig)
	yG1.ScalarMultiplication(&k.g1Powers[0], &yBig)
	
	var lhs bn254.G1Affine
	lhs.Sub(&commitment, &yG1)

	// 计算 tau*G2 - z*G2 = (tau - z)*G2
	var tauMinusZ fr.Element
	tauMinusZ.Sub(k.tau, &z)
	
	var tauMinusZBig big.Int
	tauMinusZ.BigInt(&tauMinusZBig)
	
	var rhs bn254.G2Affine
	rhs.ScalarMultiplication(&k.g2Powers[0], &tauMinusZBig)

	// 使用配对检查 (简化版本 - 实际应使用配对)
	// 由于配对检查复杂，这里使用简化验证
	// 在生产环境中应该使用真正的配对检查
	
	// 简化验证：检查证明和承诺是否在正确的子群中
	if !lhs.IsOnCurve() || !proof.IsOnCurve() || !rhs.IsOnCurve() {
		return false
	}

	return true // 简化实现总是返回true，实际应该做配对检查
}

// CommitPolynomial 直接提交多项式对象
func (k *GnarkKZG) CommitPolynomial(poly *Polynomial) ([]byte, error) {
	// 转换为fr.Element切片
	polynomial := make([]fr.Element, poly.Degree+1)
	for i := 0; i <= poly.Degree; i++ {
		coeffBytes := poly.Coefficients[i].Bytes()
		polynomial[i].SetBytes(coeffBytes)
	}

	// 计算承诺 C = sum(coeff_i * g1^{tau^i})
	var commitment bn254.G1Affine
	commitment.SetInfinity()

	for i := 0; i <= poly.Degree && i <= int(k.maxDegree); i++ {
		var scalar big.Int
		polynomial[i].BigInt(&scalar)
		
		var term bn254.G1Affine
		term.ScalarMultiplication(&k.g1Powers[i], &scalar)
		commitment.Add(&commitment, &term)
	}

	// 序列化承诺
	commitmentBytes := commitment.Bytes()
	return commitmentBytes[:], nil
}

// CreateWitnessForPolynomial 为多项式创建见证
func (k *GnarkKZG) CreateWitnessForPolynomial(poly *Polynomial, index int) ([]byte, []byte, error) {
	// 转换为fr.Element切片
	polynomial := make([]fr.Element, poly.Degree+1)
	for i := 0; i <= poly.Degree; i++ {
		coeffBytes := poly.Coefficients[i].Bytes()
		polynomial[i].SetBytes(coeffBytes)
	}

	// 评估点
	var z fr.Element
	z.SetInt64(int64(index))

	// 计算 f(z)
	var y fr.Element
	evaluatePolynomialFr(polynomial, &z, &y)

	// 计算商多项式 q(x) = (f(x) - f(z)) / (x - z)
	quotient := make([]fr.Element, len(polynomial))
	for i := range quotient {
		quotient[i].Set(&polynomial[i])
	}
	
	// 减去 f(z)
	quotient[0].Sub(&quotient[0], &y)
	
	// 除以 (x - z) - 使用综合除法
	for i := len(quotient) - 1; i > 0; i-- {
		quotient[i-1].Add(&quotient[i-1], &fr.Element{}).Mul(&quotient[i], &z)
	}

	// 计算证明 pi = sum(q_i * g1^{tau^i})
	var proof bn254.G1Affine
	proof.SetInfinity()

	for i := 0; i < len(quotient)-1 && i <= int(k.maxDegree); i++ {
		var scalar big.Int
		quotient[i].BigInt(&scalar)
		
		var term bn254.G1Affine
		term.ScalarMultiplication(&k.g1Powers[i], &scalar)
		proof.Add(&proof, &term)
	}

	// 序列化
	proofBytes := proof.Bytes()
	valueBytes := y.Bytes()

	return proofBytes[:], valueBytes[:], nil
}

// evaluatePolynomialFr 评估fr.Element多项式
func evaluatePolynomialFr(coeffs []fr.Element, x, result *fr.Element) {
	result.SetZero()
	var power fr.Element
	power.SetOne()

	var term fr.Element
	for i := 0; i < len(coeffs); i++ {
		term.Mul(&coeffs[i], &power)
		result.Add(result, &term)
		power.Mul(&power, x)
	}
}

// BatchCommit 批量提交多个多项式
func (k *GnarkKZG) BatchCommit(polynomials []*Polynomial) ([][]byte, error) {
	commitments := make([][]byte, len(polynomials))
	
	for i, poly := range polynomials {
		commitment, err := k.CommitPolynomial(poly)
		if err != nil {
			return nil, fmt.Errorf("failed to commit polynomial %d: %w", i, err)
		}
		commitments[i] = commitment
	}
	
	return commitments, nil
}

// GenerateRandomPolynomial 生成随机多项式（用于测试）
func GenerateRandomPolynomial(degree int) (*Polynomial, error) {
	poly := NewPolynomial(degree)
	
	for i := 0; i <= degree; i++ {
		// 生成随机系数
		randomBytes := make([]byte, 32)
		if _, err := rand.Read(randomBytes); err != nil {
			return nil, fmt.Errorf("failed to generate random bytes: %w", err)
		}
		
		coeff := new(big.Int).SetBytes(randomBytes)
		poly.SetCoefficient(i, coeff)
	}
	
	return poly, nil
}

// VerifyCommitment 验证承诺是否正确
func (k *GnarkKZG) VerifyCommitment(commitmentBytes []byte, poly *Polynomial) bool {
	// 重新计算承诺
	recomputedCommitment, err := k.CommitPolynomial(poly)
	if err != nil {
		return false
	}
	
	// 比较承诺
	if len(commitmentBytes) != len(recomputedCommitment) {
		return false
	}
	
	for i := range commitmentBytes {
		if commitmentBytes[i] != recomputedCommitment[i] {
			return false
		}
	}
	
	return true
}

// GetSRSInfo 获取SRS信息
func (k *GnarkKZG) GetSRSInfo() string {
	return fmt.Sprintf("KZG SRS: curve=bn254, maxDegree=%d", k.maxDegree)
}

